# 导入tushare
import tushare as ts
import pandas as pd

# 初始化pro接口
pro = ts.pro_api('854634d420c0b6aea2907030279da881519909692cf56e6f35c4718c')


def get_stock(symbol, start_date=""):
    # 拉取数据
    df = pro.daily(**{
        "ts_code": symbol,
        "trade_date": "",
        "start_date": start_date,
        "end_date": "",
        "offset": "",
        "limit": ""
    }, fields=[
        # "ts_code",
        "trade_date",
        "open",
        "high",
        "low",
        "close",
        # "pre_close",
        # "change",
        # "pct_chg",
        "vol",
        "amount"
    ])
    df.index = df['trade_date']
    del df['trade_date']
    df.rename(columns={'vol': 'volume', 'open': 'o_open', 'high': 'o_high', 'low': 'o_low', 'close': 'o_close'},
              inplace=True)

    return df


def get_factor(symbol, start_date=""):
    # 拉取数据
    df = pro.adj_factor(**{
        "ts_code": symbol,
        "trade_date": "",
        "start_date": start_date,
        "end_date": "",
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "trade_date",
        "adj_factor"
    ])
    df.index = df['trade_date']
    df.rename(columns={'adj_factor': 'factor', 'trade_date': 'date', 'ts_code':'code'}, inplace=True)
    return df


def get_symbol_data_with_factor(symbol, start_date=""):
    df = get_stock(symbol, start_date)
    df_factor = get_factor(symbol, start_date)

    df_all = pd.concat([df, df_factor], axis=1)
    for col in ['open', 'high', 'low', 'close']:
        df_all[col] = df_all['o_{}'.format(col)] * df_all['factor']
    df_all['_id'] = symbol + '_' + df_all['date']
    df_all.dropna(inplace=True)
    return df_all


if __name__ == '__main__':
    df = get_symbol_data_with_factor('600519.SH', start_date='20220821')
    print(df)
    from common.mongo_utils import write_df

    #write_df('stock_quotes', df, drop_tb_if_exist=True)
